# -*-coding:utf-8 -*-
import numpy as np
from scipy import interpolate
import matplotlib.pyplot as plt
from scipy import interpolate
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm



def func(x, y):
    return (x+y)*np.exp(-5.0*(x**2 + y**2))



# 一、一维插值
def one_dim_interpolation(x=None, y=None):
    x = np.linspace(0, 10, 11)
    #x=[  0.   1.   2.   3.   4.   5.   6.   7.   8.   9.  10.]
    y = np.sin(x)
    xnew = np.linspace(0, 10, 101)
    plt.plot(x, y, "ro")

    for kind in ["nearest","zero","slinear","quadratic","cubic"]:#插值方式
        #"nearest","zero"为阶梯插值
        #slinear 线性插值
        #"quadratic","cubic" 为2阶、3阶B样条曲线插值
        f = interpolate.interp1d(x , y, kind=kind)
        # ‘slinear’, ‘quadratic’ and ‘cubic’ refer to a spline interpolation of first, second or third order)
        ynew = f(xnew)
        plt.plot(xnew, ynew, label=str(kind))
    plt.legend(loc="lower right")
    plt.show()


# 二、二维插值
def two_dim_interpolation(x=None, y=None):

    """
    演示二维插值。
    """
    # X-Y轴分为15*15的网格
    y , x = np.mgrid[-1:1:15j, -1:1:15j]   # -1到1分成14等分 这里的虚数j只是作为一个符号，表示包含截止点的值。

    fvals = func(x,y) # 计算每个网格点上的函数值  15*15的值
    # print(len(fvals[0]))

    #三次样条二维插值
    newfunc = interpolate.interp2d(x, y, fvals, kind='cubic')

    # 计算100*100的网格上的插值
    xnew = np.linspace(-1,1,100)#x
    ynew = np.linspace(-1,1,100)#y
    fnew = newfunc(xnew, ynew)#仅仅是y值   100*100的值

    # 绘图
    # 为了更明显地比较插值前后的区别，使用关键字参数interpolation='nearest'
    # 关闭imshow()内置的插值运算。
    plt.subplot(121)
    im1 = plt.imshow(fvals, extent=[-1,1,-1,1], cmap=cm.hot, interpolation='nearest', origin="lower")#pl.cm.jet
    #extent=[-1,1,-1,1]为x,y范围  favals为
    plt.colorbar(im1)

    plt.subplot(122)
    im2 = plt.imshow(fnew, extent=[-1,1,-1,1], cmap=cm.hot, interpolation='nearest', origin="lower")
    plt.colorbar(im2)
    plt.show()


# 三、二维插值的三维展示方法

def three_dim_interpolation(x=None, y=None):
    
    """
    演示二维插值。
    """
    # X-Y轴分为20*20的网格
    x = np.linspace(-1, 1, 20)
    y = np.linspace(-1,1,20)
    x, y = np.meshgrid(x, y)#20*20的网格数据

    fvals = func(x,y) # 计算每个网格点上的函数值  15*15的值

    fig = plt.figure(figsize=(9, 6))
    #Draw sub-graph1
    ax = plt.subplot(1, 2, 1,projection = '3d')
    surf = ax.plot_surface(x, y, fvals, rstride=2, cstride=2, cmap=cm.coolwarm, linewidth=0.5, antialiased=True)
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('f(x, y)')
    plt.colorbar(surf, shrink=0.5, aspect=5)#标注

    #二维插值
    newfunc = interpolate.interp2d(x, y, fvals, kind='cubic')#newfunc为一个函数

    # 计算100*100的网格上的插值
    xnew = np.linspace(-1,1,100)#x
    ynew = np.linspace(-1,1,100)#y
    fnew = newfunc(xnew, ynew)#仅仅是y值   100*100的值  np.shape(fnew) is 100*100
    xnew, ynew = np.meshgrid(xnew, ynew)
    ax2 = plt.subplot(1, 2, 2,projection = '3d')
    surf2 = ax2.plot_surface(xnew, ynew, fnew, rstride=2, cstride=2, cmap=cm.coolwarm,linewidth=0.5, antialiased=True)
    ax2.set_xlabel('xnew')
    ax2.set_ylabel('ynew')
    ax2.set_zlabel('fnew(x, y)')
    plt.colorbar(surf2, shrink=0.5, aspect=5)#标注

    plt.show()


if __name__ == '__main__':
    three_dim_interpolation()